import pandas as pd
from matplotlib.ticker import MaxNLocator
import numpy as np
import matplotlib.pyplot as plt
import os


def dataInput(filename1):
    df = pd.read_csv('./time_cost/'+filename1)
    ret = list(df.columns)
    ret[0] = 'algorithm'
    df.columns = ret
    df = df.set_index('algorithm')
    df1 = df.T.astype(float)
    return df1


def featureDataInput(filename1):
    df = pd.read_csv('./features_number/'+filename1)
    ret = list(df.columns)
    ret[0] = 'algorithm'
    df.columns = ret
    df = df.set_index('algorithm')
    df1 = df.T.astype(int)
    return df1


def drawTime(df_draw, filename1):
    plt.figure()
    filename1 = filename1[:filename1.find('.')]
    df_draw.algorithm1.plot(ls='--', marker='v', legend=True, title=filename1)
    df_draw.algorithm2.plot(marker='o', legend=True)
    plt.savefig('./time_cost/' + filename1 + '.jpg')


def drawFeatures(df_draw, filename1):
    plt.figure()
    filename1 = filename1[:filename1.find('.')]
    df_draw.algorithm1.plot(ls='--', marker='v', legend=True, title=filename1)
    df_draw.algorithm2.plot(marker='o', legend=True)
    plt.gca().yaxis.set_major_locator(MaxNLocator(integer=True))
    plt.ylim(0, )
    plt.savefig('./features_number/'+filename1+'2.jpg')


if __name__ == '__main__':
    pathname_feature = 'features_number'
    pathname_time = 'time_cost'
    for root, dirs, files in os.walk(pathname_time):
        for filename in files:
            print(filename)
            df = dataInput(filename)
            drawTime(df, filename)

